'''
Created on Oct 26, 2009

@author: mkiyer
'''

from optparse import OptionParser
import logging
import h5py
import sys
import os
import numpy as np
import math
import operator

from veggie.sample.samplegroup import parse_samplegroups_xml
import veggie.coverage.covdb as covdb

class UTR:
    def __init__(self):
        pass
    def __str__(self):
        return '%s\t%d\t%d\t%s\t%s\t%d\t%d' % \
            (self.chrom, self.start, self.end, self.name,
             self.strand, self.polya, self.polya+1)

def read_utr_polya_file(fhd):
    '''
    symbol 
    accession
    chrom
    strand
    utr_tx_start
    utr_tx_end
    utr_cds_start
    utr_cds_end
    polya_count
    polya_starts
    polya_ends
    polya_sites
    '''
    for thisline in fhd:
        thisline = thisline.strip()
        if thisline is None:
            continue
        if thisline.startswith('#'):
            continue
        thisfields = thisline.split('\t')
        u = UTR()
        u.symbol = thisfields[0]
        u.accession = thisfields[1]
        u.chrom = thisfields[2]
        u.strand = thisfields[3]
        u.tx_start = int(thisfields[4])
        u.tx_end = int(thisfields[5])
        u.cds_start = int(thisfields[6])
        u.cds_end = int(thisfields[7])
        u.polya_count = int(thisfields[8])
        u.polya_starts = map(int, thisfields[9].split(','))
        u.polya_ends = map(int, thisfields[10].split(','))
        u.polya_sites = map(int, thisfields[11].split(','))
        yield u 

def read_polya_file(fhd):
    '''
    ACSM1   NM_052956       chr16   -       20541962        20542395        20542062
    '''
    for thisline in fhd:
        thisline = thisline.strip()
        if thisline is None:
            continue
        if thisline.startswith('#'):
            continue
        thisfields = thisline.split('\t')
        u = UTR()
        u.symbol = thisfields[0]
        u.accession = thisfields[1]
        u.chrom = thisfields[2]
        u.strand = thisfields[3]
        u.polya_count = 1
        u.polya_starts = [int(thisfields[4])]
        u.polya_ends = [int(thisfields[5])]
        u.polya_sites = [int(thisfields[6])]
        yield u

def calc_r_score(e_proximal_covs, e_distal_covs, c_proximal_covs, c_distal_covs):    
    rpkm_threshold = 1.0
    low_threshold = 0.1
    zero_func = lambda x: low_threshold if x < low_threshold else x
        
    c1 = np.median(c_proximal_covs.values())
    c2 = np.median(c_distal_covs.values())

    for sample_name in e_proximal_covs.iterkeys():
        e1 = e_proximal_covs[sample_name]
        e2 = e_distal_covs[sample_name]        
        # require some minimal RPKM within the coding region of the 
        # 3'UTR
        if c1 < rpkm_threshold or e1 < rpkm_threshold:
            yield (sample_name, 0.0)
        e1 = zero_func(e1)
        e2 = zero_func(e2)
        c1 = zero_func(c1)
        c2 = zero_func(c2)        
        r_score = math.log((e1/e2)/(c1/c2))
        yield (sample_name, r_score)

def copa_transform(covs):    
    median = np.median(covs)
    mad = np.median(np.abs(covs - median))
    if mad == 0.0:
        return np.zeros(len(covs), dtype=float)
    return (covs - median) / mad
    
def calc_copa(e_covs, c_covs):
    names = []
    covs = []
    for name, cov in e_covs.iteritems():
        names.append(name)
        covs.append(cov)
    for name, cov in c_covs.iteritems():        
        names.append(name)
        covs.append(cov)
    copas = copa_transform(np.array(covs, dtype=np.float))    
    sorted_copa = sorted(zip(names, copas), key=operator.itemgetter(1), reverse=True)
    return sorted_copa


def calc_utr_coverage(chrom, start, end, polya, strand, sample_group):
    proximal_interval = (start, polya)
    distal_interval = (polya, end) 
    # account for genes on the negative strand
    if strand == '-':
        proximal_interval, distal_interval = distal_interval, proximal_interval
    # query the coverage database for the proximal and distal coverage
    proximal_covs = {}
    for sample_name, sample_cov in sample_group.calc_avg_cov_rpkm(chrom,
                                                                  proximal_interval[0],
                                                                  proximal_interval[1]):
        proximal_covs[sample_name] = sample_cov
    distal_covs = {}
    for sample_name, sample_cov in sample_group.calc_avg_cov_rpkm(chrom,
                                                                  distal_interval[0],
                                                                  distal_interval[1]):
        distal_covs[sample_name] = sample_cov
    return proximal_covs, distal_covs

def calc_3prime_exon_coverage(chrom, start, end, sample_group):
    # get coverage of the coding part of the 3'UTR
    exon_covs = {}
    for sample_name, sample_cov in sample_group.calc_avg_cov_rpkm(chrom, start, end):
        exon_covs[sample_name] = sample_cov
    return exon_covs

def report_coverage(outfile, u, start, end, polya, e_proximal_covs, e_distal_covs, c_proximal_covs, c_distal_covs):    
    c_median_proximal_cov = np.median(c_proximal_covs.values())
    c_median_distal_cov = np.median(c_distal_covs.values())
    for sample_name in e_proximal_covs.iterkeys():
        output_fields = [u.symbol,
                         u.accession,
                         u.chrom,
                         u.strand,
                         str(start),
                         str(end),
                         str(polya),
                         sample_name,
                         '%.3f' % e_proximal_covs[sample_name],
                         '%.3f' % e_distal_covs[sample_name],
                         '%.3f' % c_median_proximal_cov,
                         '%.3f' % c_median_distal_cov]
        outfile.write('\t'.join(output_fields) + '\n')      

def report_raw_coverage(outfile, u, start, end, polya, e_proximal_covs, e_distal_covs, c_proximal_covs, c_distal_covs):
    e_proximal_str = '|'.join([','.join([n,'%.3f' % x]) for n, x in e_proximal_covs.iteritems()])
    e_distal_str = '|'.join([','.join([n,'%.3f' % x]) for n, x in e_distal_covs.iteritems()])
    c_proximal_str = '|'.join([','.join([n,'%.3f' % x]) for n, x in c_proximal_covs.iteritems()])
    c_distal_str = '|'.join([','.join([n,'%.3f' % x]) for n, x in c_distal_covs.iteritems()])
    # format and write the output
    output_fields = [u.symbol,
                     u.accession,
                     u.chrom,
                     u.strand,
                     str(start),
                     str(end),
                     str(polya),
                     e_proximal_str,
                     e_distal_str,
                     c_proximal_str,
                     c_distal_str]
    outfile.write('\t'.join(output_fields) + '\n')    

def report_copa(outfile, u, start, end, polya, e_proximal_covs, e_distal_covs, c_proximal_covs, c_distal_covs):
    # perform the copa transform
    proximal_copas = calc_copa(e_proximal_covs, c_proximal_covs)
    distal_copas = calc_copa(e_distal_covs, c_distal_covs)
    distal_copas.reverse()
    proximal_copa_str = '|'.join([','.join([n,'%.3f' % x]) for n, x in proximal_copas])
    distal_copa_str = '|'.join([','.join([n,'%.3f' % x]) for n, x in distal_copas])
    # format and write the output
    output_fields = [u.symbol,
                     u.accession,
                     u.chrom,
                     u.strand,
                     str(start),
                     str(end),
                     str(polya),
                     proximal_copa_str,
                     distal_copa_str]
    outfile.write('\t'.join(output_fields) + '\n')    


covfile_header = '\t'.join(['#symbol',
                           'accession',
                           'chrom',
                           'strand',
                           'utr_start',
                           'utr_end',
                           'polya_site',
                           'sample_name',
                           'e_proximal',
                           'e_distal',
                           'c_median_proximal',
                           'c_median_distal'])  

copafile_header = '\t'.join(['#symbol',
                             'accession',
                             'chrom',
                             'strand',
                             'utr_start',
                             'utr_end',
                             'polya_site',
                             'proximal_copas',
                             'distal_copas'])

rawfile_header = '\t'.join(['#symbol',
                            'accession',
                            'chrom',
                            'strand',
                            'utr_start',
                            'utr_end',
                            'polya_site',
                            'e_proximal_covs',
                            'e_distal_covs',
                            'c_proximal_covs',
                            'c_distal_covs'])

if __name__ == '__main__':
    logging.basicConfig(level=logging.DEBUG)
    
    # parse command line
    optionparser = OptionParser("usage: %prog [options] <utr_table.txt>")    
    optionparser.add_option("-e", "--experiment", dest="exper",
                            help="sample name of experiment")
    optionparser.add_option("-c", "--control", dest="control",
                            help="sample name of control")
    optionparser.add_option("-s", "--samples", dest="samples",
                            help="name of sample xml file")
    optionparser.add_option("--covdb", dest="covdb",
                            help="path to coverage database")
    optionparser.add_option("--name", dest="name", default="outlier",
                            help="prefix for output files from this run")
    (options, args) = optionparser.parse_args()
    
    # open coverage db
    h5file = h5py.File(options.covdb, 'r')
    # generate SampleGroupCoverage objects that pair samples with
    # coverage data
    e_sample_group = None
    c_sample_group = None
    for sgroup in parse_samplegroups_xml(options.samples):
        if sgroup.name == options.exper:
            e_sample_group = sgroup
        if sgroup.name == options.control:
            c_sample_group = sgroup
    if e_sample_group is None or c_sample_group is None:
        logging.critical("sample group not found")
        sys.exit(0)        
    e_sample_group = covdb.SampleGroupCoverage(e_sample_group, h5file)
    c_sample_group = covdb.SampleGroupCoverage(c_sample_group, h5file)

    # open output files
    covfile = open(options.name + '_coverage.txt', 'w')
    copafile = open(options.name + '_copa.txt', 'w')
    rawfile = open(options.name + '_raw_coverage.txt', 'w')
    # write information about the command-line options, etc used for this
    # run
    for outfile in (covfile, copafile):
        outfile.write('#Sample xml file: %s\n' % os.path.abspath(options.samples))
        outfile.write('#Coverage DB: %s\n' % os.path.abspath(options.covdb))
        outfile.write('#Experimental group: %s\n' % options.exper)
        outfile.write('#Control group: %s\n' % options.control)    

    # write file headers
    covfile.write('%s\n' % (covfile_header))
    copafile.write('%s\n' % (copafile_header))
    rawfile.write('%s\n' % (rawfile_header))

    lines_finished = 0
    debug_every = 100
    # read utr data
    for u in read_utr_polya_file(open(args[0])):
        lines_finished += 1
        #if lines_finished % debug_every == 0:
        logging.debug("Finished %d" % (lines_finished))

        # only consider UTRs with polya sites
        if u.polya_count == 0:
            continue
        # TODO: only consider UTRs of protein coding genes for now
        if u.cds_start == u.cds_end:
            continue    

        # each UTR may have many polyA sites, so loop through each one to
        # test for truncations
        for i in xrange(u.polya_count):
            start, end, polya = u.polya_starts[i], u.polya_ends[i], u.polya_sites[i]
            e_proximal_covs, e_distal_covs = calc_utr_coverage(u.chrom, start, end, polya, u.strand, e_sample_group)
            c_proximal_covs, c_distal_covs = calc_utr_coverage(u.chrom, start, end, polya, u.strand, c_sample_group)

            report_copa(copafile, u, start, end, polya, e_proximal_covs, e_distal_covs, c_proximal_covs, c_distal_covs)
            report_coverage(covfile, u, start, end, polya, e_proximal_covs, e_distal_covs, c_proximal_covs, c_distal_covs)    
            report_raw_coverage(rawfile, u, start, end, polya, e_proximal_covs, e_distal_covs, c_proximal_covs, c_distal_covs)

    copafile.close()
    covfile.close()
    rawfile.close()